import mne
import warnings
import os
import glob
from mne.decoding import CSP
from sklearn.pipeline import Pipeline
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.svm import SVC
from sklearn.model_selection import StratifiedKFold, cross_val_score, train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from sklearn.metrics import precision_recall_fscore_support
import numpy as np

# 抑制由GDF头部滤波元数据写反导致的读取告警
warnings.filterwarnings(
    "ignore",
    message=r"Highpass cutoff frequency .* greater than lowpass cutoff frequency .*",
    category=RuntimeWarning,
)

# 将MNE日志级别设为ERROR，仅输出错误，避免英文INFO日志
mne.set_log_level('ERROR')

# --- 1. 数据准备：加载本地 BCI Competition IV 2b 数据集 (.gdf) ---
# 目标目录：与脚本同仓库的 MI/CSP+SVM/2b/
data_dir = os.path.join(os.path.dirname(__file__), '2b')
gdf_files = sorted(glob.glob(os.path.join(data_dir, '**', '*.gdf'), recursive=True))

if len(gdf_files) == 0:
    raise FileNotFoundError(f"未在目录中找到.gdf文件: {data_dir}")

# 读取并拼接所有.gdf文件
raw_list = [mne.io.read_raw_gdf(f, preload=True, verbose=False) for f in gdf_files]
# 修正部分设备头信息中高通/低通被写反的问题，避免读取时的RuntimeWarning
for r in raw_list:
    hp = r.info.get('highpass', None)
    lp = r.info.get('lowpass', None)
    if hp is not None and lp is not None and hp > lp:
        # 交换为正确顺序：highpass 应小于 lowpass（如 0.5 < 100.）
        r.info['highpass'], r.info['lowpass'] = lp, hp
raw = mne.io.concatenate_raws(raw_list)

# 参考与滤波（MI常用8-30Hz）
raw.set_eeg_reference('average', projection=True, verbose=False)
raw.filter(8, 30, fir_design='firwin', verbose=False)

# 事件提取：2b数据集中左右手分别为 769/770（cue）
# 对于GDF，annotations的描述通常为字符串'769','770'，此处显式映射。
event_id_2b = {'769': 769, '770': 770}
events, _ = mne.events_from_annotations(raw, event_id=event_id_2b, verbose=False)

# 划分时间窗：以cue出现为0s，取0~4s想象窗口；使用基线校正至cue前
tmin, tmax = 0., 4.
epochs = mne.Epochs(raw, events, event_id=event_id_2b, tmin=tmin, tmax=tmax,
                    baseline=(0, 0), preload=True, verbose=False)

X_data = epochs.get_data()  # (n_epochs, n_channels, n_times)

# 标签映射：将事件769映射为0类（一般左手），770映射为1类（一般右手）
y_labels_raw = epochs.events[:, -1]
y_labels = np.where(y_labels_raw == 769, 0, 1)

print(f"\n--- 数据准备完成 ---")
print(f"Epochs 数据形状: {X_data.shape}")
print(f"标签形状: {y_labels.shape}")
print(f"标签类别分布（0类/1类）: {np.bincount(y_labels)}")

# --- 2. 经典方法：CSP + LDA ---
print(f"\n--- 经典方法：CSP + LDA ---")

# CSP提取4个空间模式 (2对)
# 修正：设置 reg='ledoit_wolf' 解决协方差矩阵奇异性问题，提高数值稳定性
csp = CSP(n_components=4, reg='ledoit_wolf', log=True, rank='full', transform_into='average_power')

# 构建CSP+LDA管道
clf_csp_lda = Pipeline([
    ('CSP', csp),
    ('Scaler', StandardScaler()), # 对CSP特征进行标准化
    ('LDA', LinearDiscriminantAnalysis())
])

# 交叉验证评估 (使用分层K折交叉验证)
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
# 增加 error_score='raise' 可以在再次出现失败时立即看到错误，方便调试
scores_csp_lda = cross_val_score(clf_csp_lda, X_data, y_labels, cv=cv, n_jobs=1, error_score='raise')

print(f"CSP + LDA 交叉验证准确率：")
print(f"  各折分数：{scores_csp_lda}")
print(f"  平均准确率：{scores_csp_lda.mean():.4f}")
print(f"  标准差（2倍）：+/- {scores_csp_lda.std()*2:.4f}")


# --- 3. 经典方法：CSP + SVM ---
print(f"\n--- 经典方法：CSP + SVM ---")

# 构建CSP+SVM管道
clf_csp_svm = Pipeline([
    ('CSP', csp), # 复用CSP对象
    ('Scaler', StandardScaler()), # 对CSP特征进行标准化
    ('SVM', SVC(kernel='linear', C=1.0, random_state=42)) # 使用线性核SVM
])

# 交叉验证评估
# 增加 error_score='raise'
scores_csp_svm = cross_val_score(clf_csp_svm, X_data, y_labels, cv=cv, n_jobs=1, error_score='raise')

print(f"CSP + SVM 交叉验证准确率：")
print(f"  各折分数：{scores_csp_svm}")
print(f"  平均准确率：{scores_csp_svm.mean():.4f}")
print(f"  标准差（2倍）：+/- {scores_csp_svm.std()*2:.4f}")

# 最终报告 (可选：在所有数据上训练模型并生成报告)
# 划分一次训练集和测试集，展示详细报告
X_train, X_test, y_train, y_test = train_test_split(X_data, y_labels, test_size=0.2, random_state=42, stratify=y_labels)

# 这里的 fit 也会使用正则化后的 CSP，避免 LinAlgError
clf_csp_lda.fit(X_train, y_train)
y_pred = clf_csp_lda.predict(X_test)

print("\n--- 最终评估报告（CSP + LDA，20%测试集） ---")
acc = accuracy_score(y_test, y_pred)
prec, rec, f1, _ = precision_recall_fscore_support(y_test, y_pred, labels=[0,1], zero_division=0)
cm = confusion_matrix(y_test, y_pred, labels=[0,1])
print(f"测试准确率：{acc:.4f}")
print("精确率（按类别0/1）：", np.round(prec, 4))
print("召回率（按类别0/1）：", np.round(rec, 4))
print("F1（按类别0/1）：", np.round(f1, 4))
print("混淆矩阵：\n", cm)